package main

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/sirupsen/logrus"
)

type RegistryProxy struct {
	config *Config
	client *http.Client
}

func NewRegistryProxy(config *Config) *RegistryProxy {
	return &RegistryProxy{
		config: config,
		client: &http.Client{
			Timeout: 30 * time.Second,
		},
	}
}

func (p *RegistryProxy) HandleRequest(w http.ResponseWriter, r *http.Request) {
	logrus.Infof("[HANDLER] Handling v2 request: %s %s from %s", r.Method, r.URL.Path, r.RemoteAddr)
	logrus.Debugf("[HANDLER] Full URL: %s", r.URL.String())
	logrus.Debugf("[HANDLER] User-Agent: %s", r.Header.Get("User-Agent"))
	logrus.Debugf("[HANDLER] Authorization: %s", r.Header.Get("Authorization"))

	// Parse the request path to extract repository and reference information
	path := r.URL.Path
	if strings.HasPrefix(path, "/v2/") {
		logrus.Debugf("[HANDLER] Routing to v2 handler")
		p.handleV2Request(w, r)
	} else if strings.HasPrefix(r.URL.Path, "/token") {
		// Handle authentication token requests
		logrus.Debugf("Handling authentication token request: %s", r.URL.String())
		p.handleTokenRequest(w, r)
	} else {
		logrus.Warnf("[HANDLER] Path %s does not match /v2/ prefix, returning 404", r.URL.Path)
		http.Error(w, "Not Found", http.StatusNotFound)
	}
}

// HandleGenericRequest handles all non-v2 requests, including authentication
func (p *RegistryProxy) HandleGenericRequest(w http.ResponseWriter, r *http.Request) {
	logrus.Infof("[GENERIC] Handling generic request: %s %s from %s", r.Method, r.URL.Path, r.RemoteAddr)
	logrus.Debugf("[GENERIC] Full URL: %s", r.URL.String())
	logrus.Debugf("[GENERIC] User-Agent: %s", r.Header.Get("User-Agent"))
	logrus.Debugf("[GENERIC] Authorization: %s", r.Header.Get("Authorization"))

	// Special handling for auth endpoints
	if strings.Contains(r.URL.Path, "/token") {
		logrus.Infof("[GENERIC] Detected authentication token request, routing to token handler")
		p.handleTokenRequest(w, r)
		return
	}

	// For all other requests, proxy them directly
	logrus.Debugf("[GENERIC] Proxying request with retry logic")
	p.proxyRequestWithRetry(w, r)
}

func (p *RegistryProxy) handleV2Request(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	logrus.Debugf("[V2] Processing v2 request path: %s", path)

	// Handle /v2/ endpoint (registry version check)
	if path == "/v2/" {
		logrus.Infof("[V2] Registry version check request")
		p.proxyRequestWithRetry(w, r)
		return
	}

	// Handle manifest requests
	if strings.Contains(path, "/manifests/") {
		logrus.Infof("[V2] Manifest request detected")
		p.handleManifestRequest(w, r)
		return
	}

	// Handle blob requests
	if strings.Contains(path, "/blobs/") {
		logrus.Infof("[V2] Blob request detected")
		p.handleBlobRequest(w, r)
		return
	}

	// Handle other requests by proxying
	logrus.Infof("[V2] Other v2 request, proxying directly")
	p.proxyRequestWithRetry(w, r)
}

func (p *RegistryProxy) handleManifestRequest(w http.ResponseWriter, r *http.Request) {
	logrus.Infof("[MANIFEST] Handling manifest request: %s %s", r.Method, r.URL.Path)

	// Extract repository name and tag/digest from path
	parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/v2/"), "/")
	logrus.Debugf("[MANIFEST] Path parts: %v", parts)

	if len(parts) < 3 {
		logrus.Errorf("[MANIFEST] Invalid manifest path, insufficient parts: %d", len(parts))
		http.Error(w, "Invalid manifest path", http.StatusBadRequest)
		return
	}

	// Repository name is everything except the last two parts (manifests and tag/digest)
	repoName := strings.Join(parts[:len(parts)-2], "/")
	reference := parts[len(parts)-1]

	logrus.Infof("[MANIFEST] Repository: %s, Reference: %s, Method: %s", repoName, reference, r.Method)

	// Use go-containerregistry to fetch manifest
	if r.Method == "GET" || r.Method == "HEAD" {
		logrus.Debugf("[MANIFEST] Fetching manifest with retry")
		p.fetchManifestWithRetry(w, r, repoName, reference)
	} else if r.Method == "PUT" {
		logrus.Debugf("[MANIFEST] Handling manifest push")
		// Handle manifest push
		p.handleManifestPush(w, r, repoName, reference)
	} else {
		logrus.Debugf("[MANIFEST] Proxying %s request", r.Method)
		// For other methods, proxy the request
		p.proxyRequestWithRetry(w, r)
	}
}

func (p *RegistryProxy) handleBlobRequest(w http.ResponseWriter, r *http.Request) {
	logrus.Infof("[BLOB] Handling blob request: %s %s", r.Method, r.URL.Path)
	logrus.Debugf("[BLOB] Authorization header: %s", r.Header.Get("Authorization"))

	// Extract repository name and blob digest from path
	// Path format: /v2/{name}/blobs/{digest}
	pathParts := strings.Split(strings.TrimPrefix(r.URL.Path, "/v2/"), "/")
	if len(pathParts) < 3 || pathParts[len(pathParts)-2] != "blobs" {
		logrus.Errorf("[BLOB] Invalid blob path format: %s", r.URL.Path)
		http.Error(w, "Invalid blob path", http.StatusBadRequest)
		return
	}

	repoName := strings.Join(pathParts[:len(pathParts)-2], "/")
	blobDigest := pathParts[len(pathParts)-1]

	logrus.Debugf("[BLOB] Repository: %s, Digest: %s", repoName, blobDigest)

	// Use go-containerregistry to fetch the blob
	p.fetchBlobWithGoContainerRegistry(w, r, repoName, blobDigest)
}

func (p *RegistryProxy) fetchBlobWithGoContainerRegistry(w http.ResponseWriter, r *http.Request, repoName, blobDigest string) {
	logrus.Infof("[BLOB-GCR] Fetching blob using go-containerregistry: %s@%s", repoName, blobDigest)

	// Parse digest
	digest, err := name.NewDigest(fmt.Sprintf("%s@%s", repoName, blobDigest))
	if err != nil {
		logrus.Errorf("[BLOB-GCR] Failed to parse digest %s: %v", blobDigest, err)
		http.Error(w, "Invalid digest", http.StatusBadRequest)
		return
	}

	// Use anonymous authentication for public repositories
	opts := []remote.Option{
		//remote.WithAuthFromKeychain(authn.DefaultKeychain),
	}
	auth := p.extractAuthFromRequest(r)
	var optssy []remote.Option
	if auth != nil {
		optssy = append(opts, remote.WithAuth(auth))
	} else {
		optssy = append(opts, remote.WithAuthFromKeychain(authn.DefaultKeychain))
	}

	// Fetch the manifest
	blob, err := remote.Layer(digest, optssy...)
	if err != nil {
		logrus.Errorf("[BLOB-GCR] Failed to fetch blob %s: %v", blobDigest, err)
		if strings.Contains(err.Error(), "401") || strings.Contains(err.Error(), "403") {
			w.Header().Set("Www-Authenticate", fmt.Sprintf(`Bearer realm="http://%s/token",service="registry.docker.io",scope="repository:%s:pull"`, r.Host, repoName))
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
		} else {
			http.Error(w, "Failed to fetch blob", http.StatusInternalServerError)
		}
		return
	}

	// Get blob reader
	blobReader, err := blob.Compressed()
	if err != nil {
		logrus.Errorf("[BLOB-GCR] Failed to get blob reader: %v", err)
		http.Error(w, "Failed to read blob", http.StatusInternalServerError)
		return
	}
	defer blobReader.Close()

	// Get blob size
	blobSize, err := blob.Size()
	if err != nil {
		logrus.Errorf("[BLOB-GCR] Failed to get blob size: %v", err)
		http.Error(w, "Failed to get blob size", http.StatusInternalServerError)
		return
	}

	// Set response headers
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Length", fmt.Sprintf("%d", blobSize))
	w.Header().Set("Docker-Content-Digest", blobDigest)
	w.WriteHeader(http.StatusOK)

	// Copy blob data to response
	written, err := io.Copy(w, blobReader)
	if err != nil {
		logrus.Errorf("[BLOB-GCR] Failed to copy blob data: %v", err)
		return
	}

	logrus.Infof("[BLOB-GCR] Successfully served blob %s, size: %d bytes", blobDigest, written)
}

func (p *RegistryProxy) handleTokenRequest(w http.ResponseWriter, r *http.Request) {
	logrus.Infof("[TOKEN] Handling token request: %s %s", r.Method, r.URL.Path)
	logrus.Debugf("[TOKEN] Full URL: %s", r.URL.String())

	// 构建认证服务器 URL
	authURL := "https://auth.docker.io/token"
	if r.URL.RawQuery != "" {
		authURL += "?" + r.URL.RawQuery
	}

	logrus.Debugf("[TOKEN] Proxying to auth server: %s", authURL)

	// 代理到认证服务器
	p.proxyTokenRequest(w, r, authURL)
}

func (p *RegistryProxy) fetchManifestWithRetry(w http.ResponseWriter, r *http.Request, repoName, reference string) {
	var lastErr error
	for i := 0; i <= p.config.Registry.RetryCount; i++ {
		if i > 0 {
			logrus.Debugf("Retrying manifest fetch (attempt %d/%d) for %s:%s", i+1, p.config.Registry.RetryCount+1, repoName, reference)
			time.Sleep(p.config.Registry.RetryDelay * time.Duration(i))
		}

		if p.fetchManifest(w, r, repoName, reference) {
			return // Success
		}

		lastErr = fmt.Errorf("failed to fetch manifest")
	}

	logrus.Errorf("Failed to fetch manifest after %d retries: %v", p.config.Registry.RetryCount+1, lastErr)
	p.proxyRequestWithRetry(w, r) // Fallback to proxy
}

func isPublicRepository(pathOrRepo string) bool {
	// 如果输入是URL路径，先提取仓库名称
	var repoName string
	if strings.HasPrefix(pathOrRepo, "/v2/") {
		// 从URL路径中提取仓库名称
		path := strings.TrimPrefix(pathOrRepo, "/v2/")
		parts := strings.Split(path, "/")
		if len(parts) > 0 {
			// 对于manifest和blob请求，需要找到manifests或blobs之前的部分作为仓库名
			if len(parts) >= 2 && (strings.Contains(path, "/manifests/") || strings.Contains(path, "/blobs/")) {
				for i, part := range parts {
					if part == "manifests" || part == "blobs" {
						repoName = strings.Join(parts[:i], "/")
						break
					}
				}
			} else {
				// 对于其他请求，使用第一部分作为仓库名
				repoName = parts[0]
			}
		}
	} else {
		// 直接传入的仓库名称
		repoName = pathOrRepo
	}

	// Docker Hub官方镜像没有用户名前缀，只有镜像名
	// 例如: nginx, ubuntu, mysql 等
	// 私有仓库格式: username/imagename 或 registry.com/username/imagename

	// 去除registry前缀
	parts := strings.Split(repoName, "/")

	// 如果只有一个部分，通常是官方公有镜像
	if len(parts) == 1 {
		return true
	}

	// 如果是library/开头，也是官方公有镜像
	if len(parts) == 2 && parts[0] == "library" {
		return true
	}

	return false
}

func (p *RegistryProxy) extractAuthFromRequest(r *http.Request) authn.Authenticator {
	//判断是否为公有仓库
	if isPublicRepository(r.URL.Path) {
		logrus.Debugf("[AUTH] Public repository detected for path: %s", r.URL.Path)
		return nil
	}

	auth := r.Header.Get("Authorization")
	logrus.Infof("[AUTH] Authorization header: %s", auth)

	if auth == "" {
		logrus.Debugf("[AUTH] No Authorization header found")
		return nil
	}

	// Handle Bearer token
	if strings.HasPrefix(auth, "Bearer ") {
		token := strings.TrimPrefix(auth, "Bearer ")
		logrus.Debugf("[AUTH] Using Bearer token authentication")
		return &authn.Bearer{Token: token}
	}

	// Handle Basic auth
	if strings.HasPrefix(auth, "Basic ") {
		credentials := strings.TrimPrefix(auth, "Basic ")
		decoded, err := base64.StdEncoding.DecodeString(credentials)
		if err != nil {
			logrus.Errorf("[AUTH] Failed to decode basic auth: %v", err)
			return nil
		}

		parts := strings.SplitN(string(decoded), ":", 2)
		if len(parts) != 2 {
			logrus.Errorf("[AUTH] Invalid basic auth format")
			return nil
		}

		logrus.Debugf("[AUTH] Using Basic authentication for user: %s", parts[0])
		return &authn.Basic{
			Username: parts[0],
			Password: parts[1],
		}
	}

	logrus.Debugf("[AUTH] Unknown authorization type: %s", auth)
	return nil
}

func (p *RegistryProxy) fetchManifest(w http.ResponseWriter, r *http.Request, repoName, reference string) bool {
	// Construct the full repository reference
	var ref name.Reference
	var err error

	if strings.Contains(reference, "sha256:") {
		// It's a digest
		ref, err = name.NewDigest(fmt.Sprintf("%s@%s", repoName, reference))
	} else {
		// It's a tag
		ref, err = name.NewTag(fmt.Sprintf("%s:%s", repoName, reference))
	}

	if err != nil {
		logrus.Errorf("Failed to parse reference: %v", err)
		return false
	}

	// Create remote options
	opts := []remote.Option{
		//remote.WithAuthFromKeychain(authn.DefaultKeychain),
		remote.WithContext(context.Background()),
	}
	auth := p.extractAuthFromRequest(r)

	var optssy []remote.Option
	if auth != nil {
		optssy = append(opts, remote.WithAuth(auth))
	} else {
		optssy = append(opts, remote.WithAuthFromKeychain(authn.DefaultKeychain))
	}

	// Fetch the manifest
	manifest, err := remote.Get(ref, optssy...)
	if err != nil {
		logrus.Errorf("[MANIFEST-GCR] Failed to fetch manifest: %v", err)
		return false
	}

	// Set appropriate headers
	w.Header().Set("Content-Type", string(manifest.MediaType))
	w.Header().Set("Docker-Content-Digest", manifest.Digest.String())
	w.Header().Set("Content-Length", fmt.Sprintf("%d", len(manifest.Manifest)))

	// For HEAD requests, don't write the body
	if r.Method == "HEAD" {
		w.WriteHeader(http.StatusOK)
		return true
	}

	// Write the manifest
	w.WriteHeader(http.StatusOK)
	w.Write(manifest.Manifest)
	return true
}

func (p *RegistryProxy) handleManifestPush(w http.ResponseWriter, r *http.Request, repoName, reference string) {
	logrus.Debugf("Handling manifest push for %s:%s", repoName, reference)

	// For manifest pushes, we typically want to proxy to the upstream
	// unless we're implementing a local registry cache
	p.proxyRequestWithRetry(w, r)
}

func (p *RegistryProxy) proxyTokenRequest(w http.ResponseWriter, r *http.Request, targetURL string) {
	logrus.Debugf("[TOKEN_PROXY] Starting token proxy request to: %s", targetURL)

	// 创建新的请求
	req, err := http.NewRequest(r.Method, targetURL, r.Body)
	if err != nil {
		logrus.Errorf("[TOKEN_PROXY] Failed to create request: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	// 复制请求头
	for key, values := range r.Header {
		for _, value := range values {
			req.Header.Add(key, value)
		}
	}

	// 发送请求
	resp, err := p.client.Do(req)
	if err != nil {
		logrus.Errorf("[TOKEN_PROXY] Failed to send request: %v", err)
		http.Error(w, "Bad Gateway", http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	logrus.Debugf("[TOKEN_PROXY] Received response: %d %s", resp.StatusCode, resp.Status)

	// 复制响应头
	for key, values := range resp.Header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}

	// 设置状态码
	w.WriteHeader(resp.StatusCode)

	// 复制响应体
	_, err = io.Copy(w, resp.Body)
	if err != nil {
		logrus.Errorf("[TOKEN_PROXY] Failed to copy response body: %v", err)
	}

	logrus.Debugf("[TOKEN_PROXY] Token proxy request completed successfully")
}

func (p *RegistryProxy) proxyRequestWithRetry(w http.ResponseWriter, r *http.Request) {
	logrus.Debugf("[RETRY] Starting proxy with retry logic, max retries: %d", p.config.Registry.RetryCount)

	var lastErr error
	for i := 0; i <= p.config.Registry.RetryCount; i++ {
		logrus.Debugf("[RETRY] Attempt %d/%d", i+1, p.config.Registry.RetryCount+1)

		if i > 0 {
			logrus.Warnf("[RETRY] Proxy attempt %d failed, retrying in %v...", i, p.config.Registry.RetryDelay*time.Duration(i))
			time.Sleep(p.config.Registry.RetryDelay * time.Duration(i))
		}

		if p.proxyRequest(w, r) {
			logrus.Infof("[RETRY] Proxy succeeded on attempt %d", i+1)
			return // Success
		}

		lastErr = fmt.Errorf("proxy request failed")
		logrus.Errorf("[RETRY] Proxy attempt %d failed", i+1)
	}

	logrus.Errorf("[RETRY] All %d proxy attempts failed: %v", p.config.Registry.RetryCount+1, lastErr)
	http.Error(w, "Bad Gateway", http.StatusBadGateway)
}

func (p *RegistryProxy) proxyRequest(w http.ResponseWriter, r *http.Request) bool {
	logrus.Debugf("[PROXY] Starting proxy request: %s %s", r.Method, r.URL.Path)
	logrus.Debugf("[PROXY] Original URL: %s", r.URL.String())
	logrus.Debugf("[PROXY] Request headers: %v", r.Header)

	// Create the upstream URL
	upstreamURL, err := url.Parse(p.config.Registry.Upstream)
	if err != nil {
		logrus.Errorf("[PROXY] Failed to parse upstream URL: %v", err)
		return false
	}

	// Create the proxy request
	proxyURL := *r.URL
	proxyURL.Scheme = upstreamURL.Scheme
	proxyURL.Host = upstreamURL.Host

	logrus.Infof("[PROXY] Proxying request to: %s", proxyURL.String())

	proxyReq, err := http.NewRequest(r.Method, proxyURL.String(), r.Body)
	if err != nil {
		logrus.Errorf("[PROXY] Failed to create proxy request: %v", err)
		return false
	}

	// Copy headers
	for key, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(key, value)
		}
	}

	// Add custom headers from config
	for key, value := range p.config.Registry.Headers {
		proxyReq.Header.Set(key, value)
		logrus.Debugf("[PROXY] Added custom header: %s = %s", key, value)
	}

	logrus.Debugf("[PROXY] Final request headers: %v", proxyReq.Header)

	// Make the request
	logrus.Debugf("[PROXY] Making HTTP request to upstream...")
	resp, err := p.client.Do(proxyReq)
	if err != nil {
		logrus.Errorf("[PROXY] Failed to make proxy request: %v", err)
		return false
	}
	defer resp.Body.Close()

	logrus.Infof("[PROXY] Upstream response: %d %s", resp.StatusCode, resp.Status)
	logrus.Debugf("[PROXY] Response headers: %v", resp.Header)

	// Copy response headers, but modify Www-Authenticate to point to our proxy
	for key, values := range resp.Header {
		if strings.ToLower(key) == "www-authenticate" {
			// Modify the realm in Www-Authenticate header to point to our proxy
			for _, value := range values {
				// Determine protocol based on TLS configuration
				protocol := "http"
				if p.config.Server.TlsCert != "" && p.config.Server.TlsKey != "" {
					protocol = "https"
				}
				// Replace auth.docker.io with our proxy address
				modifiedValue := strings.ReplaceAll(value, "https://auth.docker.io/token", protocol+"://"+r.Host+"/token")
				w.Header().Add(key, modifiedValue)
				logrus.Infof("[PROXY] Modified Www-Authenticate: %s -> %s", value, modifiedValue)
			}
		} else {
			for _, value := range values {
				w.Header().Add(key, value)
			}
		}
	}

	// Set status code
	w.WriteHeader(resp.StatusCode)

	// Copy response body
	if r.Method != "HEAD" {
		bodyBytes, err := io.Copy(w, resp.Body)
		if err != nil {
			logrus.Errorf("[PROXY] Failed to copy response body: %v", err)
			return false
		}
		logrus.Debugf("[PROXY] Copied %d bytes in response body", bodyBytes)
	}

	logrus.Infof("[PROXY] Request completed successfully: %s %s -> %d", r.Method, r.URL.Path, resp.StatusCode)
	return true
}